
from typing import Generator, List
from django.conf import settings
from cozepy import Coze, TokenAuth, Message, ChatEventType
from .models import FinanceChatHistory


class CozeFinanceService:
    """财务助手Coze智能体服务封装"""

    def __init__(self):
        # 初始化Coze客户端
        self.coze = Coze(
            auth=TokenAuth(token=settings.COZE_API_TOKEN2),
            base_url=settings.COZE_CN_BASE_URL,
        )
        self.bot_id = settings.COZE_FINANCE_BOT_ID  # 建议修改为财务专用的bot_id配置

    def stream_chat(self, prompt: str, user_id: str) -> Generator[str, None, None]:
        """
        流式调用Coze财务助手并记录聊天历史
        :param prompt: 用户输入的需求
        :param user_id: 唯一用户ID
        :return: 流式文本生成器
        """
        if not prompt.strip():
            raise ValueError("Prompt不能为空")

        # 收集AI完整响应内容
        full_response = []

        # 调用Coze流式接口
        for event in self.coze.chat.stream(
                bot_id=self.bot_id,
                user_id=user_id,
                additional_messages=[Message.build_user_question_text(prompt)]
        ):
            if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                content = event.message.content.strip()
                full_response.append(content)
                yield content
            elif event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                # 保存完整对话记录到数据库
                FinanceChatHistory.objects.create(
                    user_id=user_id,
                    prompt=prompt,
                    ai_response=''.join(full_response)
                )
                yield "[DONE]"